package com.town.garden.utils;

import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacpp.avcodec;
import org.bytedeco.javacpp.avutil;
import org.bytedeco.javacv.*;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.Objects;
import java.util.concurrent.CountDownLatch;

import static org.bytedeco.javacpp.avutil.AV_LOG_ERROR;

/**
 * @Description 录制rtsp流工具类
 * @E-mail cwm3412638@163.com
 * @Author cwm
 * @Date 2022/5/16
 */
@Slf4j
public class JavaCVRecordUtil {
    /**
     * 视频录制工具类 配置视频读取器参数
     *
     * @param inputFile
     * @param outputFile
     * @param audioChannel
     * @throws Exception
     * @throws org.bytedeco.javacv.FrameRecorder.Exception
     */
    public static void frameRecord(String inputFile, String outputFile, int audioChannel) {
        //只打印错误日志
        avutil.av_log_set_level(AV_LOG_ERROR);
        // 该变量建议设置为全局控制变量，用于控制录制结束
        boolean isStart = true;
        // 获取视频源
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputFile);
        //读取流超时时间
        grabber.setOption("stimoout", "5000");
        grabber.setOption("rtsp_transport", "tcp");
        grabber.setFrameRate(30);
        grabber.setVideoBitrate(3000000);
        // 流媒体输出地址，分辨率（长，高），是否录制音频（0:不录制/1:录制）
        FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(outputFile, 1280, 720, audioChannel);
        recorder.setFrameRate(30);
        recorder.setVideoBitrate(3000000);
        recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
        recorder.setFormat("mp4");
        recordByFrame(grabber, recorder, isStart);
    }

    /**
     * 视频录制器 录制视频
     *
     * @param grabber
     * @param recorder
     * @param status
     * @throws Exception
     * @throws org.bytedeco.javacv.FrameRecorder.Exception
     */
    private static void recordByFrame(FFmpegFrameGrabber grabber, FFmpegFrameRecorder recorder, Boolean status) {
        try {//建议在线程中使用该方法
            grabber.start();
            recorder.start();
            Frame frame = null;
            while (status && (frame = grabber.grabFrame()) != null) {
                recorder.record(frame);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                recorder.stop();
                recorder.release();
            } catch (FrameRecorder.Exception e) {
                e.printStackTrace();
            }
            try {
                grabber.stop();
                grabber.release();
            } catch (FrameGrabber.Exception e) {
                e.printStackTrace();
            }
        }
    }


    /**
     * rtsp 流转rtmp流 利用ffmpeg  此方法可用，但是延迟偏高新能不佳 项目中没有使用到
     *
     * @param rtmpUrl rtmp 流地址
     * @param rtspUrl rtsp 流地址
     * @param width   宽
     * @param height  高
     * @throws Exception
     */
    public static void rtspToRtmp(String rtmpUrl, String rtspUrl, int width, int height) throws Exception {
        //只打印错误日志
        avutil.av_log_set_level(AV_LOG_ERROR);
        // 使用rtsp的时候需要使用 FFmpegFrameGrabber，不能再用 FrameGrabber
        FFmpegFrameGrabber grabber = FFmpegFrameGrabber.createDefault(rtspUrl);
        // 使用tcp的方式，不然会丢包很严重
        grabber.setOption("rtsp_transport", "tcp");
        grabber.setImageWidth(width);
        grabber.setImageHeight(height);
        grabber.start();
        // 流媒体输出地址，分辨率（长，高），是否录制音频（0:不录制/1:录制）
        FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(rtmpUrl, width, height, 0);
        recorder.setInterleaved(true);
        //recorder.setVideoOption("crf","28");
        recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
        // rtmp的类型
        recorder.setFormat("flv");
        recorder.setFrameRate(25);
        recorder.setImageWidth(width);
        recorder.setImageHeight(height);
        // yuv420p
        recorder.setPixelFormat(0);
        recorder.start();
        int count = 0;
        while (true) {
            count++;
            Frame frame = grabber.grabImage();
            if (frame == null) {
                continue;
            }
            recorder.record(frame);
        }
    }

    /**
     * 截取视频帧 生成封面图
     *
     * @param filePath    文件地址
     * @param outFilePath 生成文件地址
     * @throws IOException
     */
    public static void getImage(String filePath, String outFilePath) {
        FFmpegFrameGrabber ff = null;
        try {
            File file = new File(outFilePath);
            //创建视频帧抓取工具
            ff = FFmpegFrameGrabber.createDefault(filePath);
            ff.start();
            //获取图片第一帧
            Frame frame = null;
            //获取视频的总时长
            int ftp = ff.getLengthInFrames();
            int flag = 0;
            while (flag <= ftp) {
                frame = ff.grabImage();
                //避免出现黑图
                if ((frame != null)) {
                    break;
                }
                flag++;
            }
            Java2DFrameConverter converter = new Java2DFrameConverter();
            BufferedImage bufferedImage = converter.getBufferedImage(frame);
            if (Objects.nonNull(bufferedImage)) {
                ImageIO.write(bufferedImage, "jpg", file);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (Objects.nonNull(ff)) {
                try {
                    ff.close();
                    ff.stop();
                } catch (FrameGrabber.Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
